import numpy as np
from scipy.spatial.transform import Rotation as R
from pyvox.parser import VoxParser
import os


class VoxelProcessor:
    def __init__(self, config):
        self.config = config
        self.vox_model = None
        self.voxels = None
        self.size = None
        self.center = None
        self.palette = None

    def load_vox_file(self, vox_file):
        self.vox_model = VoxParser(f"./vox-file/{vox_file}.vox").parse()
        self.voxels = self.vox_model.to_dense()

        # ✨ 交换 x 与 y 轴，修复方向错乱
        self.voxels = np.transpose(self.voxels, (0, 2, 1))  # (y, x, z)

        self.size = np.array(self.voxels.shape)
        self.center = (self.size - 1) / 2.0

        self.palette = np.array(
            [(c.r, c.g, c.b) for c in self.vox_model.palette[1:257]],
            dtype=np.uint8
        )

        if self.config.DEBUG_MODE:
            print(f"模型尺寸: {self.size}, 调色板: {self.palette.shape}")

    def generate_projection_slices(self):
        """保留原有投影功能（如需）"""
        raise NotImplementedError("未实现：投影模式逻辑")

    def generate_z_axis_y0_slices(self):
        """以Z轴旋转模型，并提取Y=0平面切片（适用于z_y0_slice模式）"""
        FRAMES = self.config.FRAMES
        SCALE = self.config.SCALE_FACTOR * 8.0 / (max(self.size[0], self.size[1]) / 2.0 + 1)

        valid_voxels = [
            (
                np.array([x, y, z]) - self.center,
                self.palette[self.voxels[x, y, z] - 1]
            )
            for x in range(self.size[0])
            for y in range(self.size[1])
            for z in range(self.size[2])
            if self.voxels[x, y, z] > 0
        ]

        frame_data = np.zeros((FRAMES, 512, 3), dtype=np.uint8)

        for frame_idx in range(FRAMES):
            theta = frame_idx * (2 * np.pi / FRAMES)
            rot_mat = R.from_euler('z', theta).as_matrix()
            slice_grid = np.zeros((16, 16, 3), dtype=np.uint8)
            depth_grid = np.full((16, 16), -np.inf)

            for pos, color in valid_voxels:
                rotated = rot_mat @ pos
                if abs(rotated[1]) < 0.5:  # 提取 y=0 层
                    x_proj = int(rotated[0] * SCALE + 8)
                    z_proj = int(rotated[2] * SCALE + 8)
                    if 0 <= x_proj < 16 and 0 <= z_proj < 16:
                        depth = rotated[1]
                        if depth > depth_grid[z_proj, x_proj]:
                            depth_grid[z_proj, x_proj] = depth
                            slice_grid[z_proj, x_proj] = color

            for y in range(32):
                for x in range(16):
                    data_x = 15 - x if self.config.FLIP_X else x
                    data_y = 15 - (y % 16) if self.config.FLIP_Y else y % 16
                    if self.config.XY_SWING:
                        data_x, data_y = data_y, data_x

                    color = slice_grid[data_y, data_x] if (0 <= data_x < 16 and 0 <= data_y < 16) else [0, 0, 0]

                    if y < 16:
                        led_x = x if y % 2 == 0 else 15 - x
                        led_index = y * 16 + led_x
                    else:
                        mirror_y = 31 - y
                        led_x = x if mirror_y % 2 == 0 else 15 - x
                        led_index = 256 + mirror_y * 16 + led_x

                    frame_data[frame_idx, led_index] = color

        return frame_data

    def generate_y0_slices(self):
        """固定底面 y=0 的切片旋转展示"""
        if self.vox_model is None:
            raise ValueError("未加载VOX文件")

        FRAMES = self.config.FRAMES
        SCALE = self.config.SCALE_FACTOR * 8.0 / (max(self.size[0], self.size[2]) / 2.0)

        frame_data = np.zeros((FRAMES, 512, 3), dtype=np.uint8)

        for frame_idx in range(FRAMES):
            theta = frame_idx * (2 * np.pi / FRAMES)
            rot_mat = R.from_euler('y', theta).as_matrix()
            slice_grid = np.zeros((16, 16, 3), dtype=np.uint8)
            depth_grid = np.full((16, 16), -np.inf) if self.config.USE_DEPTH else None

            y = 0
            for x in range(self.size[0]):
                for z in range(self.size[2]):
                    if self.voxels[x, y, z] > 0:
                        pos = np.array([x, y, z]) - self.center
                        rotated = rot_mat @ pos

                        x_proj = int(rotated[0] * SCALE + 8)
                        z_proj = int(rotated[2] * SCALE + 8)

                        if 0 <= x_proj < 16 and 0 <= z_proj < 16:
                            color = self.palette[self.voxels[x, y, z] - 1]
                            if self.config.USE_DEPTH:
                                depth = rotated[1]
                                if depth > depth_grid[z_proj, x_proj]:
                                    depth_grid[z_proj, x_proj] = depth
                                    slice_grid[z_proj, x_proj] = color
                            else:
                                slice_grid[z_proj, x_proj] = color

            for y in range(32):
                for x in range(16):
                    data_x = 15 - x if self.config.FLIP_X else x

                    if self.config.FLIP_Y:
                        data_y = 15 - y if y < 16 else 31 - y
                    else:
                        data_y = y if y < 16 else y - 16

                    if self.config.XY_SWING:
                        data_x, data_y = data_y, data_x

                    color = slice_grid[data_y, data_x] if 0 <= data_x < 16 and 0 <= data_y < 16 else [0, 0, 0]

                    if y < 16:
                        led_x = x if y % 2 == 0 else 15 - x
                        led_index = y * 16 + led_x
                    else:
                        mirror_y = 31 - y
                        led_x = x if mirror_y % 2 == 0 else 15 - x
                        led_index = 256 + mirror_y * 16 + led_x

                    frame_data[frame_idx, led_index] = color

        return frame_data

    def save_to_binary(self, frame_data, filename):
        frame_data.tofile(filename)
        file_size = os.path.getsize(filename) / 1024
        print(f"保存完成: {filename}, 大小: {file_size:.1f}KB, 帧数: {frame_data.shape[0]}")
        return file_size
